fix[next]: concat_where crashes on neighbor-list fields - #2837
fix[next]: concat_where crashes on neighbor-list fields#2837tehrengruber wants to merge 9 commits into
concat_where crashes on neighbor-list fields#2837Conversation
6e70606 to
5215204
Compare
concat_where crashes on local (neighbor-list) fieldsconcat_where crashes on neighbor-list fields
havogt
left a comment
There was a problem hiding this comment.
Review of the current head (3342813). Everything below was checked against the branch; where I ran something I say so.
The fix works. I reproduced the crash on the merge base (ee1bb4f6a) and confirmed it is gone at this head with correct numerics. test_concat_where.py is 62 passed on gtfn, mypy src/ is clean, and pre-commit passes on the changed files.
One thing that has no code anchor, so it goes here: the PR description no longer matches the diff. The description says the fix is promote_dims(domain.dims, extract_dims(tb), extract_dims(fb)) in type_synthesizer.py, but this head does not touch that file at all — the fix widens type_info.promote to accept ListType. An earlier head of this branch did implement the described approach, so the body appears to have gone stale in a force-push. Two consequences: anyone reviewing from the description reviews the wrong patch, and the "fuzz-checked over dim combinations" claim no longer applies to anything in the diff. Worth updating before merge.
Inline comments below, roughly strongest first.
| np.where(edge_mask[:, np.newaxis], a[e2v_table], b[e2v_table]), | ||
| axis=1, | ||
| initial=0, | ||
| where=e2v_table != common._DEFAULT_SKIP_VALUE, |
There was a problem hiding this comment.
This skip-value mask is dead code. E2V is constructed with skip_value=None in both mesh descriptors (cases_utils.py:300 and :395), so e2v_table != common._DEFAULT_SKIP_VALUE is all-True in every parametrization and initial=0 never applies.
So despite the uses_sparse_fields marker, these tests never exercise skip values on the local dimension. Either drop the where=/initial= as misleading, or switch to V2E, which does carry _DEFAULT_SKIP_VALUE in skip_value_mesh and would give real coverage.
(Same at :494 and :500.)
The GTIR 'concat_where' type synthesizer derives the result dims via
'type_info.promote(tb, fb)', which asserted that the promoted dtypes are
'ScalarType'. A local field carries a 'ListType' dtype at the GTIR
level, so any 'concat_where' with a local-field branch failed with a
bare 'AssertionError' during type inference -- and it did so even though
both branches have the identical dtype, since the assertion is on the
dtype's kind, not on the operands differing.
Let 'promote' handle 'ListType' the same way it handles 'ScalarType':
both promote only between equal types, so the two cases collapse into a
single check. The docstring records that a 'ListType' only ever reaches
'promote' from the ITIR level, because the frontend represents the same
concept as a field with a local dimension in 'dims' and a scalar dtype.
The combination was previously untested ('test_concat_where.py' had no
local-dimension coverage), so the latent assert never fired in CI.
3342813 to
8643e6d
Compare
`type_info.promote` accepts `ListType`s only together with other lists: element types must be equal, and an `offset_type` of `None` (a list from `make_const_list`) is compatible with any other offset, matching `map_list`. The ITIR `concat_where` type synthesizer and the frontend `where`/`concat_where` deduction now take both dims and dtype from `promote`'s result instead of checking dtype equality by hand and calling `promote` only for its dims. The frontend re-raises `promote`'s `ValueError` as a located `DSLError`, as the binary-operator deduction does. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh
The lowering wraps the non-local branch of a `concat_where` whose other branch contains a local field in `make_const_list`, as `where` already does, so a scalar branch reaches the type synthesizer as a list. `transform_to_as_fieldop` turns `concat_where` into an `if_` selecting between whole lists, and `if_` typed its result as its true branch, which drops the neighbor offset when the true branch is a constant list. For two list branches `if_` now returns their promotion. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh
The local-field `concat_where` tests borrowed `uses_sparse_fields` to be xfailed on dace, although they pass no sparse field. What dace lacks is `concat_where` with a list result, so they get a dedicated marker for that. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh
The local-field tests use `V2E` instead of `E2V`, which has skip values in `skip_value_mesh`, so the skip-value mask in their references takes effect. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh
havogt
left a comment
There was a problem hiding this comment.
lgtm, maybe @tehrengruber wants to double-check my changes once back...
There was a problem hiding this comment.
🟡 Changes recommended
Tuple-valued concat_where calls with mixed local-field and scalar leaves still fail during type inference.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Fixes concat_where type inference for neighbor-list fields.
Changes:
- Adds
ListTypepromotion and list-aware conditional inference. - Wraps scalar branches as constant lists.
- Adds local-field, scalar, tuple, and error-path tests.
File summaries
| File | Description |
|---|---|
tests/next_tests/unit_tests/type_system_tests/test_type_info.py |
Tests list promotion. |
tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py |
Updates promotion error expectation. |
tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py |
Updates where error expectation. |
tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py |
Adds tuple and neighbor-list scenarios. |
tests/next_tests/definitions.py |
Marks unsupported DaCe coverage. |
src/gt4py/next/type_system/type_info.py |
Implements list promotion. |
src/gt4py/next/iterator/type_system/type_synthesizer.py |
Promotes list-valued conditionals. |
src/gt4py/next/ffront/foast_to_gtir.py |
Promotes scalar branches to lists. |
src/gt4py/next/ffront/foast_passes/type_deduction.py |
Centralizes branch promotion and errors. |
pyproject.toml |
Registers the new test marker. |
Review details
- Files reviewed: 10/10 changed files
- Comments generated: 1
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| if not isinstance(node.type, ts.TupleType) and any( | ||
| type_info.contains_local_field(t) for t in (true_type, false_type) | ||
| ): | ||
| true_branch = promote_to_list(true_type)(true_branch) | ||
| false_branch = promote_to_list(false_type)(false_branch) |
In a tuple-valued concat_where, each leaf whose counterpart in the other branch is a local field is wrapped in make_const_list, as for non-tuple branches. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh
…ions concat_where builds one concat_where per tuple element through process_elements, as where does, promoting a non-local element to a list next to a local one. Both take the per-element path for named collections too, so a named collection holding a local field is no longer mapped or promoted as a whole. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh
…t_where Embedded where and concat_where rebuilt a NamedTuple as a plain tuple and passed a dataclass collection on to the field dispatch. Both now extract the collection, select per element and construct the same collection type again. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh
Problem
Any
concat_wherewhose branches contain a local (neighbor-list) field crashes during type inference with a bareAssertionError, e.g.:The combination was previously untested (
test_concat_where.pyhad no local-dimension coverage), so the latent assert never fired in CI.Root cause
The GTIR
concat_wheretype synthesizer derives the result dimensions viatype_info.promote(tb, fb). Besides merging the dims,promotealso promotes the branches' dtypes and asserts they are allScalarType— but at the GTIR level a local field isFieldType(dims=[Edge], dtype=ListType(...)).Fix
type_info.promoteacceptsListTypes together with other lists: their element types must be equal, and anoffset_typeofNone(a list frommake_const_list) is compatible with any other. The GTIRconcat_wheretype synthesizer and the FOAST-levelwhere/concat_wherededuction take both dims and dtype frompromote's result instead of checking dtype equality by hand; at the FOAST levelpromote'sValueErroris re-raised as a locatedDSLError, as the binary-operator deduction does.A scalar branch next to a local field, e.g.
concat_where(Edge < 2, a(E2V), 0.0), is lowered as inwhere: the non-local branch is wrapped inmake_const_list. Sincetransform_to_as_fieldopturnsconcat_whereinto anif_selecting between whole lists, theif_type synthesizer returns the promotion of two list branches instead of the true branch's type.This is a pragmatic solution to unblock #2833.